/*
 *   Copyright (C) 2009 Ultratronik Entwicklungs GmbH 
 *
 *     - Author: Sören Wellhöfer <soeren.wellhoefer@gmx.net>
 *
 *   This program is free software; you can redistribute it and/or modify
 *   it under the terms of the GNU General Public License as published by
 *   the Free Software Foundation; either version 2 of the License, or
 *   (at your option) any later version.
 *
 *
 * Compilation:
 *   gcc -o memload memload.c 
 *
 * Notes:
 *   Run as root.
 *
 * Abstract:
 *   This program is intended to stress the heap (essentially
 *   a machine's RAM) in a very particual way. It continuously
 *   reads data from the file stream given as the argument FILE upon
 *   invocation and writes it to heap memory in the following way:
 *
 *   File stream data is read in chunks and copied to a heap memory region
 *   which is permanently enlarged until the maximum memory utilization
 *   is reached (which can be set using -memory p%). It defaults to 100%.
 *   Then the memory is deallocated chunk by chunk until all is free'd.
 *   All caches and buffers will get dropped (flushed) all along also. 
 *   Then this entire process starts all over again until somebody decides 
 *   to kill the program. I call this method `swingfill'.
 *
 *   Generally, this all is done to do create a situation of severe
 *   memory use with lots of writes and removes for benchmarking purposes.
 *
 *   Optionally, when `-chunk n' is given, it is attempted to read
 *   up to n bytes at once. This number defaults to 512. Blocking io is used.
 *
 *   By using the `-fixed' option it is possible to fix the system memory
 *   utilization at a certain percentage. No clearing or freeing of
 *   allocated memory space will be performed.
 *
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <stdarg.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/sysinfo.h>

/* Constants */
#define DEFAULT_CHUNK_BYTES 512
#define DEFAULT_MAX_MEMORY_USAGE 100
#define FIXED_SLEEP_PERIOD 2000 * 1000 /* 0.5sec */

const char *usg = "%s [-chunk b] [-memory p%] [-fixed] [-verbose] FILE\n";
const char *prog;

/* Macros */
#define TO_KB(BYTE) (((double)BYTE) / 1024.0)

/* Flags */
enum {
	F_VERBOSE = (unsigned int)(1<<0),
	F_FIXED = (unsigned int)(1<<1)
};
unsigned int flags;

/* Globals */
char **memory_list;
size_t chunk_size;
FILE *in_stream;
int max_mem;

/* Prototypes */
static void* memalloc(size_t);

/* Implementation */
static inline void usage()
{
	printf(usg, prog);
}

static void errmsg(char *frm, int with_errno, ...)
{
	static const char *err = "error";
	static const char *format = "%s: %s\n";
	size_t len;
	char *pre;
	va_list l;

	va_start(l, with_errno);

	len = strlen(err) + strlen(format) + strlen(frm);
	pre = (char*)memalloc((len + 1) * sizeof(char));
	pre[len] = '\0';
	
	sprintf(pre, format, err, frm);
	vfprintf(stderr, pre, l);

	if(with_errno)
		fprintf(stderr, "(%s)\n", strerror(errno));

	free(pre);
	va_end(l);
}

static double memutil()
{
	struct sysinfo sysi;

	if(sysinfo(&sysi) < 0) {
		errmsg("unable to obtain memory information", 1);
		exit(-1);
	}
	
	return (double)(sysi.totalram - sysi.freeram - sysi.bufferram) / (double)sysi.totalram * 100.0;
}

const char *dead =
	"\n      .-  \n"
	"     (O.O)  \n"
	"   '=.|m|.=' \n"
	"   .='`\"``=.\n\n"
	" Out of Memory\n\n";

static void* memalloc(size_t s)
{
	void *mem = (void*)malloc(s);
	if(!mem) {
		write(1, dead, strlen(dead));
		exit(-1);
	}
	return mem;
}

static void* memrealloc(void *m, size_t s)
{
	void *mem = (void*)realloc(m, s);
	if(s > 0 && !mem) write(1, dead, strlen(dead));
	return mem;
}

static void drop_caches()
{
	int fd = open("/proc/sys/vm/drop_caches", O_WRONLY);
	sync();
	write(fd, "3", 1);
	close(fd);
}

static void free_mem(int size, int newsize)
{
	while(--size >= newsize-1) {
		free(memory_list[size]);
		memory_list[size] = NULL;
		
		if(flags & F_VERBOSE)
			printf("free_chunk=%d, realloc_size=%.2f kB\n",
			       size, TO_KB(size * chunk_size));
	}
	
	memory_list = (char**)memrealloc(memory_list, chunk_size * newsize);
	drop_caches();
	if(!newsize) memory_list = NULL;
}

static void fill()
{
	int memsize, read;
	char *mp;

	memory_list = NULL;
	mp = NULL;
	
	for(memsize = 1; ;drop_caches()) {

		if(flags & F_VERBOSE)
			printf("mem_use=%.2f%%, mem_chunks=%d, mem_size=%.2f kB\n", memutil(),
			       memsize-1, TO_KB(chunk_size * (memsize-1)));

		if(memutil() >= (double)max_mem) {
			memsize--;
			if(flags & F_FIXED) {
				read = memsize;
				memsize *= max_mem / memutil();
				if(read < 0) read = memsize = 0;
				free_mem(read, memsize);
				usleep(FIXED_SLEEP_PERIOD);
				if(!memsize) memsize = 1;
				continue;
			} else {
				free_mem(memsize, 0);
				memsize = 1;
			}
		} else {
			memory_list = (char**)memrealloc(memory_list, memsize * chunk_size);
			mp = memory_list[memsize-1] = (char*)memalloc(chunk_size);
			
			if(!mp) goto error;
			read = fread(mp, sizeof(char), chunk_size, in_stream);

			if(ferror(in_stream)) {
				errmsg("erroneous input stream", 0);
				goto error;
			}

			/* Fill holes */
			memset(mp + read, ' ', chunk_size - read);
			memsize++;
		}
	}
	
 error:
	while(--memsize >= 0) free(memory_list[memsize]);
	free(memory_list);
}

int main(int argc, char **argv)
{
	prog = *argv++;
	
	if(argc < 2) {
		usage();
		return -1;
	}
	
	chunk_size = DEFAULT_CHUNK_BYTES;
	max_mem = DEFAULT_MAX_MEMORY_USAGE;
	in_stream = NULL;
	flags = 0;
	
	if(geteuid() != 0)
		fprintf(stderr, "warning: not running as root\n");
	
	for( ; --argc > 0; argv++) {

		if(!strcmp("-chunk", *argv)) {
			
			if((chunk_size = atoi(*++argv)) == 0) {
				errmsg("not a valid chunk size: %s", 0, *argv);
				return -1;
			}

			argc--;
			
		} else if(!strcmp("-memory", *argv)) {

			max_mem = atoi(*++argv);

			if(max_mem < 0 || max_mem > 100) {
				errmsg("not a valid percentage (0..100)", 0, *argv);
				return -1;
			}

			argc--;
		} else if(!strcmp("-verbose", *argv)) {

			flags |= F_VERBOSE;
		} else if(!strcmp("-fixed", *argv)) {

			flags |= F_FIXED;
		}
		else {
			
			if((in_stream = fopen(*argv, "r")) == NULL) {
				errmsg("unable to open %s for reading", 1, *argv);
				return -1;
			}
		}
	}

	if(!in_stream) {
		errmsg("no input file stream", 0);
		return -1;
	}

	double mem_used = memutil();
	
	if(mem_used >= (double)max_mem) {
		errmsg("maximum memory use already reached: %d%% >= %d%%", 0, (int)mem_used, (int)max_mem);
		return -1;
	}

	fill();
	
	if(in_stream) free(in_stream);

	return 0;
}
